import torch
import torch.nn as nn
import torch.utils.data
import numpy as np
import VegetableDataset
import util

def prepareData(model:nn.Module,dataset:VegetableDataset.VegetableDataset,save_path:str):
    dataloader=torch.utils.data.DataLoader(dataset,batch_size=util.Parameter.BATCH_SIZE)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    this_model=model.to(device)
    this_model.eval()
    all_features = []
    all_labels=[]
    for inputs,labels in dataloader:
        inputs=inputs.to(device)
        outputs=this_model(inputs).detach().cpu().numpy()
        all_features.append(outputs)
        all_labels.append(labels.numpy())
    features_np=np.vstack(all_features)
    labels_np=np.hstack(all_labels)
    np.savez_compressed(save_path, features=features_np, labels=labels_np)

if __name__=="__main__":

    model=torch.load("model/prepared/prepared_resnet50.pt",weights_only=False)
    model.fc=nn.Identity()
    train_dataset=VegetableDataset.VegetableDataset("train.csv",VegetableDataset.transform)
    validate_dataset=VegetableDataset.VegetableDataset("validate.csv",VegetableDataset.transform)
    test_dataset=VegetableDataset.VegetableDataset("test.csv",VegetableDataset.transform)
    prepareData(model,train_dataset,"train.npz")
    prepareData(model,validate_dataset,"validate.npz")
    prepareData(model,test_dataset,"test.npz")